from langchain_community.embeddings import OllamaEmbeddings
from langchain.llms import Ollama
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document
from config import OLLAMA_BASE_URL, LLM_MODEL

def main():
    # 1. 定义一些示例文档
    docs = [
        Document(page_content="LangChain是一个用于开发由语言模型驱动的应用程序的框架。"),
        Document(page_content="Ollama允许用户在本地运行大型语言模型，如Llama 3。" ),
        Document(page_content="FAISS是Facebook AI开发的一个用于高效相似性搜索和密集向量聚类的库。" ),
        Document(page_content="RAG（检索增强生成）结合了检索模型和生成模型，以提供更准确、更有依据的回答。" ),
    ]

    # 2. 初始化Ollama嵌入模型和LLM
    # 确保Ollama服务正在运行，并且已经拉取了nomic-embed-text和llama3模型
    # ollama pull nomic-embed-text
    # ollama pull llama3

    llm = Ollama(base_url=OLLAMA_BASE_URL, model=LLM_MODEL)
    embeddings = OllamaEmbeddings(base_url=OLLAMA_BASE_URL, model=LLM_MODEL)
    # 3. 使用FAISS创建向量存储
    # from_documents方法会处理文本的嵌入和索引
    print("正在创建FAISS向量存储...")
    vectorstore = FAISS.from_documents(docs, embeddings)
    print("向量存储创建成功。")

    # 4. 创建一个检索器
    retriever = vectorstore.as_retriever()

    # 在此处测试retriever
    test_question = "什么是LangChain？"
    print(f"\n正在使用retriever检索问题: '{test_question}'")
    retrieved_docs = retriever.invoke(test_question)
    print("\n检索到的内容块:")
    for i, doc in enumerate(retrieved_docs):
        print(f"  文档 {i+1}:\n    内容: {doc.page_content}")

    # 5. 定义一个提示模板
    template = """
    请根据以下上下文来回答问题。
    上下文: {context}
    问题: {question}
    回答: 
    """
    prompt = PromptTemplate.from_template(template)

    # 6. 创建RAG链
    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    # 7. 提出问题并获取答案
    question = "什么是LangChain？"
    print(f"\n正在查询: {question}")
    response = rag_chain.invoke(question)
    print(f"回答: {response}")

    question = "Ollama有什么用？"
    print(f"\n正在查询: {question}")
    response = rag_chain.invoke(question)
    print(f"回答: {response}")

if __name__ == "__main__":
    main()